'''
Author: kkk 997610780@qq.com
Date: 2025-03-05 14:41:57
LastEditors: kkk 997610780@qq.com
LastEditTime: 2025-03-06 11:53:42
FilePath: \get-video-images\get-video.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
import cv2
from PIL import Image
import numpy as np
import os
import uuid

def extract_frames(video_source):
    """
    从视频中提取帧并保存为图片
    
    参数:
        video_source: 可以是视频文件路径(字符串)或视频对象(cv2.VideoCapture)
    
    返回:
        保存图片的文件夹路径
    """
    # 创建固定的frames文件夹
    frames_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "frames")
    if not os.path.exists(frames_dir):
        os.makedirs(frames_dir)
        print(f"创建主文件夹: {frames_dir}")
    
    # 处理视频源
    if isinstance(video_source, str):
        # 如果是字符串，认为是视频路径
        video_path = video_source
        cap = cv2.VideoCapture(video_path)
        # 获取视频文件名（不含路径和扩展名）
        video_name = os.path.splitext(os.path.basename(video_path))[0]
    else:
        # 如果不是字符串，认为是视频对象
        cap = video_source
        video_name = "video"  # 默认名称
    
    # 在frames文件夹中创建以视频名称为前缀的随机命名子文件夹
    random_folder_name = f"{video_name}_{uuid.uuid4().hex[:8]}"
    images_dir = os.path.join(frames_dir, random_folder_name)
    if not os.path.exists(images_dir):
        os.makedirs(images_dir)
        print(f"创建子文件夹: {images_dir}")
    
    isOpened = cap.isOpened()  # 判断是否打开
    if not isOpened:
        print("无法打开视频")
        return None
    
    # 视频信息获取
    fps = cap.get(cv2.CAP_PROP_FPS)
    print(f"视频帧率: {fps} fps")
    
    imageNum = 0
    sum = 0
    # 修改为每秒截取一帧，计算需要跳过的帧数
    timef = int(fps)  # 每隔fps帧（即1秒）保存一张图片
    
    while isOpened:
        sum += 1
        (frameState, frame) = cap.read()  # 记录每帧及获取状态
        
        if frameState == True and (sum % timef == 0):
            # 格式转变，BGRtoRGB
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            # 转变成Image
            frame = Image.fromarray(np.uint8(frame))
            
            frame = np.array(frame)
            
            # RGBtoBGR满足opencv显示格式
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            
            imageNum = imageNum + 1
            fileName = os.path.join(images_dir, f'image{imageNum}.jpg')  # 修改为保存到子文件夹
            cv2.imwrite(fileName, frame, [cv2.IMWRITE_JPEG_QUALITY, 100])
            print(f"{fileName} successfully write in")  # 输出存储状态
            
        elif frameState == False:
            break
    
    print('finish!')
    # 如果是我们打开的视频，需要释放
    if isinstance(video_source, str):
        cap.release()
    
    return images_dir